{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0 0.0\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "from TMM_functions import PQ_matrices as pq\n",
    "from TMM_functions import scatter_matrices as sm\n",
    "from TMM_functions import redheffer_star as rs\n",
    "from TMM_functions import generate_initial_conditions as ic\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from numpy.linalg import cond\n",
    "import cmath;\n",
    "from scipy import linalg as LA\n",
    "from numpy.linalg import solve as bslash\n",
    "import time\n",
    "from convolution_matrices.convmat1D import *\n",
    "# from RCWA_1D_functions.grating_fft.grating_conv import *\n",
    "\n",
    "# Moharam et. al Formulation for stable and efficient implementation for RCWA\n",
    "plt.close(\"all\")\n",
    "'''\n",
    "'''\n",
    "L0 = 1e-6;\n",
    "e0 = 8.854e-12;\n",
    "mu0 = 4*np.pi*1e-8;\n",
    "\n",
    "n_ridge = 3.48;              # ridge\n",
    "n_groove = 3.48;                # groove (unit-less)\n",
    "lattice_constant = 0.7;  # SI units\n",
    "\n",
    "## simulation parameters\n",
    "theta = (0)*np.pi/180;\n",
    "phi = (0)*np.pi/180;\n",
    "e_r = 1;\n",
    "m_r = 1;\n",
    "e_t = 1;\n",
    "m_t = 1;\n",
    "n_i = np.sqrt(e_r);\n",
    "\n",
    "normal_vector= [0,0,-1]\n",
    "\n",
    "# epsilon_tensor = np.array([[12,1,0],[0,2,0],[1,0,12]]);\n",
    "# epsilon_tensor = np.array([[12,10,0],[10,12,0],[0,0,2]]);\n",
    "# epsilon_tensor = np.array([[12,5,0],[5,12,0],[0,0,12]]);\n",
    "# epsilon_tensor = np.array([[12,1,0],[0,2,0],[1,0,12]]);\n",
    "\n",
    "epsilon_tensor = np.array([[12,0,10],[0,12,0],[10,0,12]]);\n",
    "## if the exy eyx term is equal to exx eyy, then there is a degeneracy which you have to account for.\n",
    "epsilon_tensor = np.array([[12,0,0],[0,12,5],[0,10,12]]);\n",
    "\n",
    "mu_tensor = np.array([[1,0,0],[0,1,0],[0,0,1]])\n",
    "I = np.matrix(np.eye(2))\n",
    "\n",
    "## specify the initial kx and ky\n",
    "k0=1\n",
    "kx = k0 *np.sin(theta)*np.cos(phi);\n",
    "ky = k0 *np.sin(theta)*np.sin(phi);\n",
    "print(kx, ky)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## redefined functions \n",
    "def P(kx, ky, e_r, m_r):\n",
    "        return (1/e_r)*np.matrix([[kx * ky - exy,  m_r * e_r - kx ** 2],\n",
    "                                  [ky ** 2 - m_r * e_r,  -kx * ky -exy]]);\n",
    "\n",
    "def Q(kx, ky, e_r, m_r):\n",
    "        Q = (1/m_r)*np.matrix([[kx * ky, (m_r * e_r) - kx ** 2],\n",
    "                               [ky ** 2 - (m_r * e_r), -kx * ky]]);\n",
    "        return Q;\n",
    "\n",
    "def P_Q_kz(kx, ky, e_r, mu_r):\n",
    "    '''\n",
    "    r is for relative so do not put epsilon_0 or mu_0 here\n",
    "    :param kx:\n",
    "    :param ky:\n",
    "    :param e_r:\n",
    "    :param mu_r:\n",
    "    :return:\n",
    "    '''\n",
    "    # no eigensolution needed in homogeneous medium.\n",
    "    kz = np.sqrt(mu_r * e_r - kx ** 2 - ky ** 2);\n",
    "    q = Q(kx, ky, e_r, mu_r)\n",
    "    p = P(kx, ky, e_r, mu_r)\n",
    "\n",
    "    return p, q, kz;    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'polarization_amplitudes' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-40937a88347a>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[1;31m## ========================================== ##\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 23\u001b[1;33m \u001b[1;33m[\u001b[0m\u001b[0mpte\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mptm\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mpolarization_amplitudes\u001b[0m\u001b[1;33m;\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     24\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     25\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mwavelengths\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m  \u001b[1;31m# in SI units\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'polarization_amplitudes' is not defined"
     ]
    }
   ],
   "source": [
    "ref = [];\n",
    "trans = [];\n",
    "I = np.matrix(np.eye(2, 2));  # unit 2x2 matrix\n",
    "\n",
    "[e_r, m_r] = [1,1];\n",
    "[e_t, m_t] = [1,1];\n",
    "n_i = np.sqrt(e_r*m_r);\n",
    "kx = n_i * np.sin(theta) * np.cos(phi);  # constant in ALL LAYERS; kx = 0 for normal incidence\n",
    "ky = n_i * np.sin(theta) * np.sin(phi);  # constant in ALL LAYERS; ky = 0 for normal incidence\n",
    "normal_vector = np.array([0, 0, -1])  # positive z points down;\n",
    "ate_vector = np.matrix([0, 1, 0]);  # vector for the out of plane E-field\n",
    "\n",
    "## =================  specify gap media ========================##\n",
    "e_h = 1; m_h = 1;\n",
    "Pg, Qg, kzg = pq.P_Q_kz(kx, ky, e_h, m_h)\n",
    "Wg = I;  # Wg should be the eigenmodes of the E field, which paparently is the identity, yes for a homogeneous medium\n",
    "sqrt_lambda = cmath.sqrt(-1) * Wg;\n",
    "# remember Vg is really Qg*(Omg)^-1; Vg is the eigenmodes of the H fields\n",
    "Vg = Qg * Wg * (sqrt_lambda) ** -1;\n",
    "\n",
    "## ========================================== ##\n",
    "\n",
    "[pte, ptm] = [1,0];\n",
    "\n",
    "for i in range(len(wavelengths)):  # in SI units\n",
    "    ## initialize global scattering matrix: should be a 4x4 identity so when we start the redheffer star, we get I*SR\n",
    "\n",
    "    Sg11 = np.matrix(np.zeros((2, 2)));\n",
    "    Sg12 = np.matrix(np.eye(2, 2));\n",
    "    Sg21 = np.matrix(np.eye(2, 2));\n",
    "    Sg22 = np.matrix(np.zeros((2, 2)));  # matrices\n",
    "    Sg = np.block(\n",
    "        [[Sg11, Sg12], [Sg21, Sg22]]);  # initialization is equivelant as that for S_reflection side matrix\n",
    "\n",
    "    ### ================= Working on the Reflection Side =========== ##\n",
    "    Pr, Qr, kzr = pq.P_Q_kz(kx, ky, e_r, m_r)\n",
    "\n",
    "    ## ============== values to keep track of =======================##\n",
    "    S_matrices = list();\n",
    "    kz_storage = [kzr];\n",
    "    X_storage = list();\n",
    "    ## ==============================================================##\n",
    "\n",
    "    # define vacuum wavevector k0\n",
    "    lam0 = wavelengths[i];  # k0 and lam0 are related by 2*pi/lam0 = k0\n",
    "    k0 = 2*np.pi/lam0;\n",
    "    ## modes of the layer\n",
    "    Om_r = np.matrix(cmath.sqrt(-1) * kzr * I);\n",
    "    X_storage.append(Om_r);\n",
    "    W_ref = I;\n",
    "    V_ref = Qr * Om_r.I;  # can't play games with V like with W because matrices for V are complex\n",
    "\n",
    "    ## calculating A and B matrices for scattering matrix\n",
    "    Ar, Br = sm.A_B_matrices(Wg, W_ref, Vg, V_ref);\n",
    "\n",
    "    S_ref, Sr_dict = sm.S_R(Ar, Br);  # scatter matrix for the reflection region\n",
    "    S_matrices.append(S_ref);\n",
    "    Sg, D_r, F_r = rs.RedhefferStar(Sg, S_ref);\n",
    "\n",
    "    ## go through the (ANISOTROPIC) layers\n",
    "    for i in range(len(ER)):\n",
    "        # ith layer material parameters\n",
    "        e = ER[i];\n",
    "        m = UR[i];\n",
    "\n",
    "        # longitudinal k_vector\n",
    "        P, Q, kzl = pq.P_Q_kz(kx, ky, e, m)\n",
    "\n",
    "        ## E-field modes that can propagate in the medium\n",
    "        W_i = I;\n",
    "        ## corresponding H-field modes.\n",
    "        Om = cmath.sqrt(-1) * kzl * I;\n",
    "        X_storage.append(Om)\n",
    "        V_i = Q * np.linalg.inv(Om);\n",
    "\n",
    "        # now defIne A and B\n",
    "        A, B = sm.A_B_matrices(Wg, W_i, Vg, V_i);\n",
    "\n",
    "        # calculate scattering matrix\n",
    "        S_layer, Sl_dict = sm.S_layer(A, B, layer_thicknesses[i], k0, Om)\n",
    "        S_matrices.append(S_layer);\n",
    "\n",
    "        ## update global scattering matrix using redheffer star\n",
    "        Sg, D_i, F_i = rs.RedhefferStar(Sg, S_layer);\n",
    "\n",
    "    ##========= Working on the Transmission Side==============##\n",
    "    Pt, Qt, kz_trans = pq.P_Q_kz(kx, ky, e_t, m_t);\n",
    "    kz_storage.append(kz_trans);\n",
    "\n",
    "    Om = cmath.sqrt(-1) * kz_trans * I;\n",
    "    Vt = Qt * np.linalg.inv(Om);\n",
    "\n",
    "    # get At, Bt\n",
    "    At, Bt = sm.A_B_matrices(Wg, I, Vg, Vt)\n",
    "\n",
    "    ST, ST_dict = sm.S_T(At, Bt)\n",
    "    S_matrices.append(ST);\n",
    "    # update global scattering matrix\n",
    "    Sg, D_t, F_t = rs.RedhefferStar(Sg, ST);\n",
    "\n",
    "    K_inc_vector = n_i * k0 * np.matrix([np.sin(theta) * np.cos(phi), \\\n",
    "                                         np.sin(theta) * np.sin(phi), np.cos(theta)]);\n",
    "\n",
    "    # cinc is the c1+\n",
    "    E_inc, cinc, Polarization = ic.initial_conditions(K_inc_vector, theta, normal_vector, pte, ptm)\n",
    "\n",
    "    ## COMPUTE FIELDS\n",
    "    Er = Sg[0:2, 0:2] * cinc;  # S11; #(cinc = initial mode amplitudes), cout = Sg*cinc; #2d because Ex, Ey...\n",
    "    Et = Sg[2:, 0:2] * cinc;  # S21\n",
    "\n",
    "    Er = np.squeeze(np.asarray(Er));\n",
    "    Et = np.squeeze(np.asarray(Et));\n",
    "\n",
    "    Erx = Er[0];\n",
    "    Ery = Er[1];\n",
    "    Etx = Et[0];\n",
    "    Ety = Et[1];\n",
    "\n",
    "    # apply the grad(E) = 0 equation to get z components\n",
    "    Erz = -(kx * Erx + ky * Ery) / kzr;\n",
    "    Etz = -(kx * Etx + ky * Ety) / kz_trans;  ## using divergence of E equation here\n",
    "\n",
    "    # add in the Erz component to vectors\n",
    "    Er = np.matrix([Erx, Ery, Erz]);  # a vector\n",
    "    Et = np.matrix([Etx, Ety, Etz]);\n",
    "\n",
    "    R = np.linalg.norm(Er) ** 2;\n",
    "    T = np.linalg.norm(Et) ** 2;\n",
    "    ref.append(R);\n",
    "    trans.append(T);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
